import { ActionSpec } from '../Actuators/ActionSpec';
import { IActuator } from '../Actuators/IActuator';
import { IBuiltInActuator } from '../Actuators/IBuiltInActuator';
import { BuiltInActuatorType } from '../Actuators/IBuiltInActuator';
import { BuiltInSensorType, IBuiltInSensor } from '../Sensors/IBuiltInSensor';
import { ISensor } from '../Sensors/ISensor';

/**
 * 推理（inference）事件
 */
export class InferenceEvent {
	behaviorName: string = null!;
	barracudaModelSource: string = null!;
	barracudaModelVersion: string = null!;
	barracudaModelProducer: string = null!;
	barracudaPackageVersion: string = null!;
	inferenceDevice: number = 0;
	observationSpecs: Array<EventObservationSpec> = null!;
	actionSpec: EventActionSpec = null!;
	actuatorInfos: Array<EventActuatorInfo> = null!;
	memorySize: number = 0;
	totalWeightSizeBytes: number = 0;
	modelHash: string = null!;
}


export class EventObservationSpec {
	sensorName: string = null!;

	compressionType: string = null!;

	builtInSensorType: number = 0;

	observationType: number = 0;

	dimensionInfos: EventObservationDimensionInfo[] = null!;

	public static fromSensor(sensor: ISensor): EventObservationSpec {
		const obsSpec = sensor.getObservationSpec();
		const shape = obsSpec.shape;
		const dimProps = obsSpec.dimensionProperties;
		const dimInfos: EventObservationDimensionInfo[] = new Array(shape.length);
		for (let i = 0; i < shape.length; i++) {
			dimInfos[i] = {
				size: shape[i],
				flags: dimProps[i]
			};
		}

		const builtInSensorType = (sensor as unknown as IBuiltInSensor)?.getBuiltInSensorType() ?? BuiltInSensorType.Unknown;

		return {
			sensorName: sensor.getName(),
			compressionType: sensor.getCompressionSpec().sensorCompressionType.toString(),
			builtInSensorType: builtInSensorType,
			observationType: obsSpec.observationType,
			dimensionInfos: dimInfos
		};
	}

}

export class EventActionSpec {
	numContinuousActions: number = 0;
	numDiscreteActions: number = 0;
	branchSizes: number[] = null!;

	static fromActionSpec(actionSpec: ActionSpec): EventActionSpec {
		const branchSizes = actionSpec.branchSizes ?? new Array<number>();
		return {
			numContinuousActions: actionSpec.numContinuousActions,
			numDiscreteActions: actionSpec.numDiscreteActions,
			branchSizes: branchSizes,
		};
	}
}

export class EventActuatorInfo {
	builtInActuatorType: number = 0;
	numContinuousActions: number = 0;
	numDiscreteActions: number = 0;

	static fromActuator(actuator: IActuator): EventActuatorInfo {
		let builtInActuatorType: BuiltInActuatorType = BuiltInActuatorType.Unknown;

		if ('getBuiltInActuatorType' in actuator) {
			builtInActuatorType = (actuator as IBuiltInActuator).getBuiltInActuatorType();
		}

		const actionSpec = actuator.actionSpec;

		return {
			builtInActuatorType: builtInActuatorType,
			numContinuousActions: actionSpec.numContinuousActions,
			numDiscreteActions: actionSpec.numDiscreteActions,
		};
	}
}

/**
 * 关于观测的一个维度信息
 */
export class EventObservationDimensionInfo {
	size: number = 0;
	flags: number = 0;
}


export class TrainingEnvironmentInitializedEvent {
	trainingSessionGuid: string = null!;
	trainerPythonVersion: string = null!;
	mLAgentsVersion: string = null!;
	mLAgentsEnvsVersion: string = null!;
	torchVersion: string = null!;
	torchDeviceType: string = null!;
	numEnvironments: number = 0;
	numEnvironmentParameters: number = 0;
	runOptions: string = null!;
}

export enum RewardSignals {
	None = 0,
	// 外部奖励
	Extrinsic = 1 << 0,
	// GAIL
	Gail = 1 << 1,
	// 好奇心
	Curiosity = 1 << 2,
	// RND
	Rnd = 1 << 3,
}

export enum TrainingFeatures {
	None = 0,
	// 行为克隆
	BehavioralCloning = 1 << 0,
	// 循环
	Recurrent = 1 << 1,
	// 多线程
	Threaded = 1 << 2,
	// 自我对弈
	SelfPlay = 1 << 3,
	// 课程
	Curriculum = 1 << 4,
}

export class TrainingBehaviorInitializedEvent {

	trainingSessionGuid?: string = null!;

	behaviorName: string = null!;
	trainerType: string = null!;
	rewardSignalFlags: RewardSignals = RewardSignals.Extrinsic;
	trainingFeatureFlags: TrainingFeatures = TrainingFeatures.BehavioralCloning;
	visualEncoder: string = null!;
	numNetworkLayers: number = 0;
	numNetworkHiddenUnits: number = 0;
	config: string = null!;
}

export class  RemotePolicyInitializedEvent {
	public trainingSessionGuid: string = null!;
	public behaviorName: string = null!;
	public observationSpecs: EventObservationSpec[] = [];
	public actionSpec: EventActionSpec = null!;
	public actuatorInfos: EventActuatorInfo[] = [];
	public mlAgentsEnvsVersion: string = null!;
	public trainerCommunicationVersion: string = null!;
}

